# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""custom callback
This sample code is applicable to Ascend, CPU and GPU.
"""
from mindspore import save_checkpoint
from mindspore.train.callback import Callback


class SaveCallback(Callback):
    """SaveCallback"""

    def __init__(self, eval_model, ds_eval):
        """init"""
        super(SaveCallback, self).__init__()
        self.model = eval_model
        self.ds_eval = ds_eval
        self.acc = 0

    def step_end(self, run_context):
        """step end"""
        cb_params = run_context.original_args()
        step_num = cb_params.cur_step_num
        if step_num % 10 == 0:
            result = self.model.eval(self.ds_eval, dataset_sink_mode=False)
            if result['Accuracy'] > self.acc:
                self.acc = result['Accuracy']
                print("=" * 80)
                print("ACC:: " + str(self.acc) + "\n")
                file_name = str("ARN_ucf_CROSS") + str(self.acc) + ".ckpt"
                save_checkpoint(save_obj=cb_params.train_network,
                                ckpt_file_name=file_name)
                print("Save the maximum accuracy checkpoint, the accuracy is", self.acc)
